package vip.sone.util;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.*;
import java.math.BigInteger;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * RSA加解密工具类
 *
 * @author SOneWinstone jianwenzhen@qq.com
 */
public class RSAUtil {
    private static final String RSA = "RSA";
    /**
     * RSA最大加密明文大小
     *  key的modulus是1024bit的，也就是128byte，根据RSA加密规则，加密1 byte字节的数据，需要12 byte，即其他11byte可能用于记录其他信息什么的，
     *  而1024bit长度的key则最多可以加密128-11=117byte的数据
     */
    private static final int MAX_ENCRYPT_BLOCK = 117;

    /**
     * 随机生成RSA密钥对(默认秘钥长度为1024)
     *
     * @return 密钥对
     */
    public static KeyPair generateRSAKeyPair() {
        return generateRSAKeyPair(1024);
    }

    /**
     * 随机生成RSA密钥对
     *
     * @param keyLength 秘钥长度，范围：512～2048
     * @return 密钥对
     */
    public static KeyPair generateRSAKeyPair(int keyLength) {
        try {
            KeyPairGenerator generator = KeyPairGenerator.getInstance(RSA);
            generator.initialize(keyLength);
            return generator.generateKeyPair();
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 加密数据
     *
     * @param data 原始数据
     * @param publicKey 公钥
     * @return 加密数据
     */
    public static byte[] encryptData(byte[] data, PublicKey publicKey) {
        try {
            // Java默认"RSA"="RSA/ECB/PKCS1Padding"
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            // 编码前设定编码方式及公钥
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            // 传入编码数据及返回编码结果
            int inputLen = data.length;
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            int offset = 0;
            byte[] cache;
            // 对数据分段加密(最大长度117)
            while (inputLen - offset > 0) {
                if (inputLen - offset > MAX_ENCRYPT_BLOCK) {
                    cache = cipher.doFinal(data, offset, MAX_ENCRYPT_BLOCK);
                } else {
                    cache = cipher.doFinal(data, offset, inputLen - offset);
                }
                out.write(cache, 0, cache.length);
                offset += MAX_ENCRYPT_BLOCK;
            }
            byte[] encryptData = out.toByteArray();
            out.close();
            return encryptData;
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | BadPaddingException | InvalidKeyException | IOException | IllegalBlockSizeException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 解密数据
     *
     * @param data 加密数据
     * @param privateKey 私钥
     * @return 解密后的数据
     */
    public static byte[] decryptData(byte[] data, PrivateKey privateKey) {
        try {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            return cipher.doFinal(data);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 通过公钥byte[]还原公钥
     *
     * @param keyBytes 公钥byte[]
     * @return 公钥
     * @throws NoSuchAlgorithmException KeyFactory.getInstance() 的参数算法找不到
     * @throws InvalidKeySpecException 无效的KeySpec
     */
    public static PublicKey restorePublicKey(byte[] keyBytes) throws NoSuchAlgorithmException, InvalidKeySpecException {
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory factory = KeyFactory.getInstance(RSA);
        return factory.generatePublic(keySpec);
    }

    /**
     * 通过数据流还原公钥
     *
     * @param in 输入流
     * @return 公钥
     */
    public static PublicKey restorePublicKey(InputStream in) {
        try {
            return restorePublicKey(readKey(in));
        } catch (IOException e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 通过私钥byte[]还原私钥
     *
     * @param keyBytes 私钥byte[]
     * @return 私钥
     * @throws NoSuchAlgorithmException KeyFactory.getInstance() 的参数算法找不到
     * @throws InvalidKeySpecException 无效的KeySpec
     */
    public static PrivateKey restorePrivateKey(byte[] keyBytes) throws NoSuchAlgorithmException, InvalidKeySpecException {
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory factory = KeyFactory.getInstance(RSA);
        return factory.generatePrivate(keySpec);
    }

    /**
     * 还原字符串私钥
     *
     * @param privateKeyStr 私钥字符串
     * @return 私钥
     */
    public static PrivateKey restorePrivateKey(String privateKeyStr) {
        byte[] bytes = Base64.getDecoder().decode(privateKeyStr);
        try {
            return restorePrivateKey(bytes);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 还原字符串公钥
     *
     * @param publicKeyStr 公钥字符串
     * @return 公钥
     */
    public static PublicKey restorePublicKey(String publicKeyStr) {
        byte[] bytes = Base64.getDecoder().decode(publicKeyStr);
        try {
            return restorePublicKey(bytes);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 使用N、e值还原公钥
     */
    public static PublicKey restorePublicKey(String modulus, String publicExponent)
            throws NoSuchAlgorithmException, InvalidKeySpecException
    {
        BigInteger bigIntModulus = new BigInteger(modulus);
        BigInteger bigIntPrivateExponent = new BigInteger(publicExponent);
        RSAPublicKeySpec keySpec = new RSAPublicKeySpec(bigIntModulus, bigIntPrivateExponent);
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        return keyFactory.generatePublic(keySpec);
    }

    /**
     * 使用N、d值还原私钥
     */
    public static PrivateKey restorePrivateKey(String modulus, String privateExponent)
            throws NoSuchAlgorithmException, InvalidKeySpecException
    {
        BigInteger bigIntModulus = new BigInteger(modulus);
        BigInteger bigIntPrivateExponent = new BigInteger(privateExponent);
        RSAPublicKeySpec keySpec = new RSAPublicKeySpec(bigIntModulus, bigIntPrivateExponent);
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * 从流中读入秘钥信息
     *
     * @param in 输入流
     * @return 秘钥
     */
    private static String readKey(InputStream in) throws IOException {
        BufferedReader reader = new BufferedReader(new InputStreamReader(in));
        String line = null;
        StringBuilder builder = new StringBuilder();
        while ((line = reader.readLine()) != null) {
            if (line.charAt(0) != '-') {
                builder.append(line);
            }
        }
        return builder.toString();
    }

    /**
     * 打印公钥信息
     */
    public static void printPublicKeyInfo(PublicKey publicKey) {
        RSAPublicKey rsaPublicKey = (RSAPublicKey) publicKey;
        System.out.println("----------RSAPublicKey----------");
        System.out.println("Modulus.length=" + rsaPublicKey.getModulus().bitLength());
        System.out.println("Modulus=" + rsaPublicKey.getModulus().toString());
        System.out.println("PublicExponent.length=" + rsaPublicKey.getPublicExponent().bitLength());
        System.out.println("PublicExponent=" + rsaPublicKey.getPublicExponent().toString());
    }

    public static void printPrivateKeyInfo(PrivateKey privateKey) {
        RSAPrivateKey rsaPrivateKey = (RSAPrivateKey) privateKey;
        System.out.println("----------RSAPrivateKey ----------");
        System.out.println("Modulus.length=" + rsaPrivateKey.getModulus().bitLength());
        System.out.println("Modulus=" + rsaPrivateKey.getModulus().toString());
        System.out.println("PrivateExponent.length=" + rsaPrivateKey.getPrivateExponent().bitLength());
        System.out.println("PrivateExponent=" + rsaPrivateKey.getPrivateExponent().toString());

    }
}
